# -*- coding: utf-8 -*-
# @Time    : 2021/11/3 15:46
# @Author  : huangwei
# @File    : server2.py
# @Software: PyCharm
import base64
import os
import cv2
import numpy as np
from flask import request, Flask
from wsgiref.simple_server import make_server
import json
from method import create_dir
import time
from config_args import args
from excel_method import split_floors, get_table_box_data, write_excel
from text import TextDetector, TextRecognizer
from method import get_lines, get_text_boxes, roll_image_horize, sort_line, alter_close_lines, fix_up_lines, line_cross, \
    alter_lines, fill_rect_table
from draw_method import draw_table
from table_line_net import table_net

# 动态分配内存
import tensorflow as tf

config = tf.compat.v1.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.5
config.gpu_options.allow_growth = True
session = tf.compat.v1.InteractiveSession(config=config)

# 检测文本框和识别文字方法
text_detector = TextDetector(args)
text_recognizer = TextRecognizer(args)

# 加载识别表格线的模型
table_line_model_path = 'det_table_line_models/table-line.h5'
table_line_model = table_net((None, None, 3), 2)
table_line_model.load_weights(table_line_model_path)

app = Flask(__name__)


# 图片转excel
@app.route("/img2excel", methods=['POST'])
def img2excel():
    upload_file = request.get_data()
    req = json.loads(upload_file)

    if upload_file:
        filepath = req['filepath']
        print("upload filepath is:", filepath)
        filename = os.path.basename(filepath)
        sub_filename = os.path.splitext(filename)[0]

        img_str = req['image']
        img_decode = base64.b64decode(img_str)  # 解base64编码，得图片的二进制
        img_np_ = np.frombuffer(img_decode, np.uint8)
        img = cv2.imdecode(img_np_, cv2.COLOR_RGB2BGR)  # 转为opencv格式

        # 1. 写入图片到本地
        temp_dir = "./temp_path/{}".format(sub_filename)
        create_dir(temp_dir)
        img_path = "{}/{}".format(temp_dir, filename)
        cv2.imencode('.png', img)[1].tofile(img_path)

        # 2. 图片进行旋转调正 # 5s
        roll_path = "{}/roll_{}".format(temp_dir, filename)
        roll_image_horize(img_path, roll_path, table_line_model)
        # 不进行旋转调整操作
        # roll_path = img_path

        # 3. 找到所有的横线和竖线，去除掉不平行的和连接断续的 # 0.3s
        input_size = (1024, 1024)
        row_lines, col_lines = get_lines(roll_path, table_line_model, input_size)

        # 4. 调整为水平或者竖直
        row_lines = alter_lines(row_lines)
        col_lines = alter_lines(col_lines, axis=1)

        # 5. 两条竖线或两条横线贴近的去除一条,多执行几次,有的线段需要合并多次
        for i in range(2):
            row_lines = alter_close_lines(row_lines)
            col_lines = alter_close_lines(col_lines, axis=1)

        # 6. 横线竖线相交处处理
        row_lines, col_lines = line_cross(row_lines, col_lines)

        # 7. 对线进行从左到右从上到下排序
        row_lines = sort_line(row_lines)
        col_lines = sort_line(col_lines, axis=1)

        # 8. 识别文本框 # 16s
        # 在这一步添加上识别处文字的部分,然后将文字和框一一对应
        det_text_boxes, det_word = get_text_boxes(img, text_detector, text_recognizer, row_lines)

        print(len(det_text_boxes), len(det_word))

        # 9. 依次扫描每一条横线和竖线，判断其是否需要去除冒头的部分或延长冒头的部分。 # 0.5s
        # 先处理横线或先处理竖线可能会少了部分线段
        # 可以两次处理后叠加,再合并所有的横线竖线
        row_lines, col_lines = fix_up_lines(row_lines, col_lines, det_text_boxes, roll_path)

        # 10. 补全最外围的线,即确保表格为矩形
        row_lines, col_lines = fill_rect_table(row_lines, col_lines)

        # 11. 两条竖线或两条横线贴近的去除一条,主要是合并补全的外围的线,再进行排序
        row_lines = alter_close_lines(row_lines)
        col_lines = alter_close_lines(col_lines, axis=1)
        row_lines = sort_line(row_lines)
        col_lines = sort_line(col_lines, axis=1)

        # 12. 画出识别出的线识别出的表格
        table_box_path = "{}/table_box_{}".format(temp_dir, filename)
        table_boxes = draw_table(img, row_lines + col_lines, table_box_path)

        # 13. 划分分成多少行和多少列,便于对应进excel中
        x_index, y_index, row_lines, col_lines = split_floors(row_lines, col_lines)

        # 14. 获取表中提取的信息
        box_infos = get_table_box_data(roll_path, text_recognizer, table_boxes, det_text_boxes, det_word, x_index,
                                       y_index)

        # 15. 写入excel
        excel_path = "{}/{}.xlsx".format(temp_dir, sub_filename)
        write_excel(box_infos, x_index, y_index, excel_path)

        res = json.dumps(box_infos)

        return res


if __name__ == "__main__":
    try:
        server = make_server("", port=6075, app=app)
        print("Running on http://localhost:6075/")
        server.serve_forever()
    except:
        pass
